'''
Author: Aurora 2074375758@qq.com
Date: 2022-04-20 13:15:17
LastEditTime: 2024-02-19 15:30:36
FilePath: /Cat vs Dog/utils/dataset.py
Description: 数据集加载
Copyright (c) 2024 by Aurora, All Rights Reserved. 
'''

from PIL import Image
from torchvision import transforms
from torch.utils import data
import os
import matplotlib.pyplot as plt
import numpy as np


class Mydata(data.Dataset):
    """定义自己的数据集"""

    def __init__(self, root, Transforms=None, train=True):
        """
        Args:
            root:训练集的路径
            Transforms:图片处理的方式
            train:是否是训练数据。True:训练数据;False:验证数据
        """
        imgs = [os.path.join(root, img) for img in os.listdir(root)]
        #按序号排序，以便训练集和验证集不会混乱
        # imgs=sorted(imgs,key=lambda x:int((x.split("/")[-1]).split("_")[0]))
        imgs_num = len(imgs)
        """进行数据集的划分"""
        if train:
            self.imgs = imgs[:int(0.8 * imgs_num)]  #80%训练集
        else:
            self.imgs = imgs[int(0.8 * imgs_num):]  #20%验证集
        if Transforms is None:
            # Pytorch图像预处理时，通常使用transforms.Normalize(mean, std)对图像按通道进行标准化，
            # 即减去均值，再除以方差。这样做可以加快模型的收敛速度。其中参数mean和std分别表示图像每个通道的均值和方差序列。
            # Imagenet数据集的均值和方差为：mean=(0.485, 0.456, 0.406)，std=(0.229, 0.224, 0.225)，
            # 因为这是在百万张图像上计算而得的，所以我们通常见到在训练过程中使用它们做标准化。而对于特定的数据集，选择这个值的结果可能并不理想
            normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            # 图片处理 brightness: 亮度调整因子
            color_aug = transforms.ColorJitter(brightness=0.1)
            # transforms.Compose()类用来串联多个图片变换操作，对图像进行各种转换操作，使用Compose将这些转换操作组合起来。
            self.transforms = transforms.Compose(
                [transforms.CenterCrop([224, 224]),  #数据转换操作，用于从图像中心裁剪图片
                 transforms.Resize([224, 224]),
                 color_aug,
                 transforms.RandomHorizontalFlip(p=0.5),  # 数据增强 随机水平翻转给定的PIL.Image,概率为0.5。即：一半的概率翻转，一半的概率不翻转
                 transforms.RandomVerticalFlip(p=0.5),  # 随机旋转
                 transforms.ColorJitter(brightness=0.1),  # 亮度变化
                 transforms.ToTensor(),
                 # transforms.ToTensor()是‌PyTorch中的一个函数，主要用于将‌PIL Image或‌numpy.ndarray格式的图像转换为‌Torch.Tensor数据类型。这种转换是图像数据处理中的一个重要步骤，因为它允许模型更有效地处理这些数据
                 normalize])

    def __getitem__(self, index):
        """
		返回一张图片的数据
		训练集和验证集,则对应的是dog返回1,猫则返回0
		"""
        img_path = self.imgs[index]

        label = 1 if 'dog' in img_path.split('/')[-1] else 0  #获得标签
        #print(img_path)
        data = Image.open(img_path).convert("RGB")  #可能有单通道的，注意转换为三通道图片
        try:
            data = self.transforms(data)  #图片处理
        except:
            print(img_path)
            raise ValueError("图片打开失败")
        return data, label

    def __len__(self):
        """返回数据集中所有图片的个数"""
        return len(self.imgs)

    def getall(self):
        return self.imgs


if __name__ == "__main__":
    root = "../data/train"
    train = Mydata(root, train=True)
    img, label = train.__getitem__(4)
    imgs = train.getall()
    img = img.numpy()
    print(type(img))
    print(img.shape, label)
    print(len(train))
    # 主要是Pytorch中使用的数据格式与plt.imshow()函数的格式不一致
    # Pytorch中为[Channels, H, W]
    # 而plt.imshow()中则是[H, W, Channels]
    np_image = np.transpose(img, [1, 2, 0])*0.5 + 0.5 # 转置后做逆归一化
    plt.imshow(np_image)
    plt.show()